import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
from mindspore import Parameter, Tensor

from typing import Type

class MLPBlock(nn.Cell):
    def __init__(
        self,
        embedding_dim: int,
        mlp_dim: int,
        act: Type[nn.Cell] = nn.GELU,
    ) -> None:
        super().__init__()
        self.lin1 = nn.Dense(embedding_dim, mlp_dim)
        self.lin2 = nn.Dense(mlp_dim, embedding_dim)
        self.act = act()
    
    def construct(self, x: ms.Tensor) -> ms.Tensor:
        return self.lin2(self.act(self.lin1(x)))

# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa
# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119  # noqa
class LayerNorm2d(nn.Cell):
    def __init__(
        self, 
        num_channels: int,
        eps: float = 1e-6
    ) -> None:
        super().__init__()
        self.weight = Parameter(ops.ones(num_channels, ms.float32))
        self.bias = Parameter(ops.zeros(num_channels, ms.float32))
        self.eps = eps
    
    def construct(self, x: ms.Tensor) -> ms.Tensor:
        u = x.mean(1, keep_dims=True)
        s = ops.pow(x - u, 2).mean(1, keep_dims=True)
        x = (x - u) / ops.sqrt(s + self.eps)
        x = self.weight[:, None, None] * x + self.bias[:, None, None]
        return x

if __name__ == '__main__':
    import numpy as np
    data = np.random.randn(1, 256)
    np.save(r'D:\2023Q1\SAM-ms\data.npy', data)
    
    a = Tensor(data)
    model = LayerNorm2d(256)
    c = model(a)
    print(c.shape)